#!/usr/bin/env python2
'''
HTTP/WebSocket server serving to control the Mars rover.

Serves static content from http directory, and websocket connection
on /mars.
'''
import datetime
import socket
import sys

import tornado.httpclient
import tornado.ioloop
import tornado.web
import tornado.websocket

import mars_data

MESSAGES = {
    u'connected': {
        'en': u'Connected to Mars rover.',
        'fi': u'Yhteys Mars-m\xf6nkij\xe4\xe4n luotu.',
    },
    u'connection lost': {
        'en': u'Connection to Mars was lost, reconnecting...',
        'fi': u'Yhteys Marsiin katkesi, korjataan...',
    },
    u'connection failed': {
        'en': u'Connecting to Mars failed, retrying...',
        'fi': u'Yhteydenotto Marsiin ei onnistunut, kokeillaan uudestaan...',
    },
    u'chosen': {
        'en': u'You are now in control of the Mars rover.',
        'fi': u'Olet nyt Mars-m\xf6nkij\xe4n ohjaaja.',
    },
    u'not chosen': {
        'en': u'Someone else is currently controlling the Mars rover.',
        'fi': u'Joku muu on t\xe4ll\xe4 hetkell\xe4 ohjaamassa Mars-m\xf6nkij\xe4\xe4.',
    },
    u'protocol error': {
        'en': u'Protocol error.',
        'fi': u'Protokollavirhe.',
    },
    u'admin ok': {
        'en': u'Administrative privileges granted.',
        'fi': u'Yll\xe4pit\xe4j\xe4n oikeudet annettu.',
    },
    u'admin fail': {
        'en': u'You are not worthy of administrative power.',
        'fi': u'Et ole yll\xe4pit\xe4j\xe4n arvoinen.',
    },
    u'admin needed': {
        'en': u'You need to be administrator for this action.',
        'fi': u'Vain yll\xe4pit\xe4j\xe4 saa tehd\xe4 noin.',
    },
    u'data saved': {
        'en': u'Data saved.',
        'fi': u'Tiedot tallennettu.',
    },
    u'halt': {
        'en': u'Control of the Mars rover has been halted.',
        'fi': u'Mars-m\xf6nkij\xe4 on pys\xe4ytetty.',
    },
    u'unhalt': {
        'en': u'Control of the Mars rover has been restored.',
        'fi': u'Mars-m\xf6nkij\xe4\xe4 saa taas ohjata.',
    },
    u'bad ping': {
        'en': u'Got pong for no ping.',
        'fi': u'Pongi tuli olemattomalle pingille.',
    },
}

class WebSocketClientManager(object):
    '''
    Manages connections from multiple clients to prevent two clients from
    trying to control the Mars rover at the same time.
    '''
    ADMIN_PASSWORD = u'admin' # might want to change this
    
    def __init__(self):
        self._clients = set()
        self._admins = set()
        self.current = None
    
    def register_client(self, client):
        '''
        New client connected.
        '''
        self._clients.add(client)
    
    def request_admin(self, client, password):
        '''
        Request promotion from client to admin.
        '''
        if password == self.ADMIN_PASSWORD:
            self._clients.discard(client)
            self._admins.add(client)
            self.unicast(u'admin ok', client)
            if client is self.current:
                self.change_control()
            return True
        self.unicast(u'admin fail', client)
        return False
    
    def request_control(self, client):
        '''
        Client requesting control of the Mars rover.
        '''
        if self.current in (None, client) or client in self._admins:
            self.current = client
            self.unicast(u'chosen')
        else:
            self.unicast(u'not chosen', client)
        if not raspberry.active:
            raspberry.connect()
    
    def unregister_client(self, client):
        '''
        Some client disconnected.
        '''
        self._clients.discard(client)
        self._admins.discard(client)
        if client is self.current:
            self.change_control(tell=False)
    
    def change_control(self, client=None, tell=True):
        '''
        Change which client is in control.
        '''
        old_client = self.current
        if client:
            self.current = client
        elif len(self._clients) == 0:
            self.current = None
        elif len(self._clients) == 1 or self.current is None:
            self.current = iter(self._clients).next()
        else:
            chosen = None
            choose = False
            for client in self._clients:
                if client is self.current:
                    choose = True
                elif choose:
                    chosen = client
                    break
            self.current = chosen or iter(self._clients).next()
        self.unicast(u'chosen')
        if tell and old_client is not self.current:
            self.unicast(u'not chosen', old_client)
    
    def broadcast(self, key):
        '''
        Broadcast a text message to all connected clients.
        '''
        text = MESSAGES.get(key, {'en': key})
        for client in self._clients|self._admins:
            client.write_message(
                u'TEXT:%s:%s' % (key, text.get(client.language))
            )
    
    def unicast(self, key, client=None):
        '''
        Send a text message to specific client.
        '''
        if client is None:
            client = self.current
        if client is not None:
            text = MESSAGES.get(key, {'en': key})
            client.write_message(
                u'TEXT:%s:%s' % (key, text.get(client.language))
            )
    
    def unicast_error(self, errinst, client=None):
        '''
        Send an error message to specific client.
        '''
        if client is None:
            client = self.current
        if client is not None:
            client.write_message(
                u'TEXT:error:%s' % (errinst,)
            )
    
    def send_all(self, data):
        '''
        Send data to all connected clients.
        '''
        textdata = data.decode('utf-8', 'replace')
        for client in self._clients|self._admins:
            client.write_message(textdata)
    
    def send(self, data, client=None):
        '''
        Send data to specific client.
        '''
        textdata = data.decode('utf-8', 'replace')
        if client is None:
            client = self.current
        if client is not None:
            client.write_message(textdata)


ws = WebSocketClientManager()

class RPiConnection(object):
    '''
    Direct socket connections and webcam connection to Raspberry Pi.
    '''
    RPI_ADDRESS = ('192.168.1.213', 7000)
    VIDEO_SOURCE = (RPI_ADDRESS[0], 8080, '/?action=stream')
    RESTORE_TIME = 1 # time in seconds to try reconnecting after error
    RECONNECT_TIME = 10 # time in seconds to try reconnecting if can't connect
    TIMEOUT = 0.5 # timeout for connection attempts
    
    def __init__(self, ioloop):
        self.socket = None
        self.stream = None
        self.active = False
        self.ioloop = ioloop
        self.video_client = None
        self.video_feeds = set()
        self._halted = False
    
    def connect(self):
        '''
        Attempt to connect to Raspberry Pi.
        '''
        print "RPi connect()"
        if self.socket is not None:
            self.close()
        try:
            self.socket = socket.create_connection(self.RPI_ADDRESS,
                timeout = self.TIMEOUT)
        except socket.error as error:
            print "Error."
            print str(error)
            self.close()
            ws.broadcast(u'connection failed')
            self.reconnect(self.RECONNECT_TIME)
        else:
            # register the connection with ioloop
            self.active = True
            ws.broadcast(u'connected')
            self.stream = tornado.iostream.IOStream(self.socket,
                io_loop=self.ioloop)
            self.stream.set_close_callback(self._stream_closed)
            self.stream.read_until('\n', self._message) # start reading
            self.start_video()
    
    def start_video(self):
        '''
        Request for video data.
        '''
        print "Trying to start video"
        self.video_client = MJPEG_Proxy(
            ioloop=self.ioloop)
        self.video_client.fetch(
            self.VIDEO_SOURCE,
            frame_callback = self._got_video,
            connect_timeout = self.TIMEOUT,
        )
        print "done with start_video"
    
    def _got_video(self, data):
        '''
        Got stream of video, stream it to clients.
        '''
        # TODO: if possible, do video editing ;D
        # print "Got video! %d"%len(data)
        for client in tuple(self.video_feeds):
            client.got_video(data)
    
    def add_video_feed(self, client):
        '''
        Add a client as receiver of video feed.
        '''
        self.video_feeds.add(client)
    
    def remove_video_feed(self, client):
        '''
        Remove a client as receiver of video feed.
        '''
        self.video_feeds.remove(client)
    
    def reconnect(self, reconnect_time):
        '''
        Schedule a reconnection.
        '''
        print "Scheduling a reconnection to RPi"
        self.active = True
        when = datetime.timedelta(seconds=reconnect_time)
        self.ioloop.add_timeout(when, self.connect)
    
    def _stream_closed(self):
        '''
        Stream was closed for some reason.
        '''
        print "RPi _stream_closed"
        ws.broadcast(u'connection lost')
        self.stream = None
        if self.socket is not None:
            self.close()
        self.reconnect(self.RESTORE_TIME)
            
    def close(self):
        '''
        Close connection.
        '''
        if self.stream is not None:
            self.socket = None
            self.stream.close()
        elif self.socket is not None:
            self.socket.close()
            self.socket = None
        if self.video_client is not None:
            self.video_client.close()
        self.active = False
    
    def _message(self, message):
        '''
        Callback from IOStream, finished reading a message from RPi.
        '''
        message = message.strip()
        if message:
            ws.send_all('mars:'+message)
            self.stream.read_until('\n', self._message) # wait for new message
    
    def send(self, message):
        '''
        Send a message to RPi.
        '''
        if self.stream is not None and not self._halted:
            self.stream.write(message)

    def halt(self):
        '''
        Halt controls.
        '''
        if self.stream is not None:
            self.stream.write('s\n')
        self._halted = True
    
    def unhalt(self):
        '''
        Unhalt controls.
        '''
        self._halted = False


class MJPEG_Proxy(object):
    def __init__(self, ioloop):
        self.ioloop = ioloop
        self.socket = None
        self.stream = None
        self._callback = None
        self.boundary = None

    def close(self):
        if self.stream is not None:
            self.stream.close()
        self.__init__(self.ioloop)
    
    def fetch(self, dest,
        frame_callback=None, connect_timeout=None):
        '''
        Fetch mjpeg stream and stream it on.
        '''
        try:
            self.socket = socket.create_connection(dest[0:2],
                timeout = connect_timeout)
        except socket.error:
            return False
        else:
            # register the connection with ioloop
            self.stream = tornado.iostream.IOStream(self.socket,
                io_loop=self.ioloop)
            self.stream.set_close_callback(self._stream_closed)
            # write request
            self.stream.write('''GET %s HTTP/1.1\r
Host: %s:%d\r
Connection: close\r
Accept: */*\r
User-Agent: Hurrikaani\r
\r
''' % (
                dest[2], dest[0], dest[1],
            ))
            self.stream.read_until('\r\n\r\n', self._headers) # start reading
            self._callback = frame_callback
            return True
    
    def _stream_closed(self):
        '''
        Stream closed?
        '''
        self.__init__(self.ioloop)
    
    def _headers(self, data):
        '''
        Got headers.
        '''
        lines = data.split('\r\n')[1:-2]
        headers = dict(line.split(': ', 1) for line in lines)
        ct = headers['Content-Type']
        self.boundary = ct.split('boundary=', 1)[1]
        self.stream.read_until('--%s\r\n'%(self.boundary,), self._frame)
    
    def _frame(self, data):
        '''
        Got frame.
        '''
        endlength = len(self.boundary)+6
        if len(data) > endlength:
            headers, frame = data.split('\r\n\r\n', 1)
            if self._callback:
                self._callback(frame[:-endlength])
        self.stream.read_until('--%s\r\n'%(self.boundary,), self._frame)

        
class MarsWebSocket(tornado.websocket.WebSocketHandler):
    '''
    WebSocket handler to relay communications between RPi and web browser.
    '''
    PING_INTERVAL = datetime.timedelta(seconds=10)
    def open(self):
        '''
        Called by Tornado when opening new connection.
        '''
        self.language = 'en'
        self.fails = 0
        self.admin = False
        self.connected = datetime.datetime.utcnow()
        self.pings = {}
        ws.register_client(self)
        ioloop.add_timeout(self.PING_INTERVAL, self.send_ping)

    def write_message(self, data):
        '''
        Error-handling wrapper around WebSocketHandler.write_message.
        '''
        try:
            super(MarsWebSocket, self).write_message(data)
        except (AttributeError, tornado.websocket.WebSocketClosedError):
            ws.unregister_client(self)

    def on_message(self, message):
        '''
        Called by Tornado when new message is received.
        '''
        params = message.split(u':')
        try:
            if params[0] == u'set':
                if len(params) != 3:
                    ws.unicast(u'protocol error', self)
                elif params[1] == u'language':
                    self.language = params[2]
                elif params[1] in mars_data.SUPPORTED_DATA:
                    if self.admin:
                        mars_data.save_data(params[1], params[2])
                        ws.unicast(u'data saved', self)
                    else:
                        ws.unicast(u'admin needed', self)
                else:
                    ws.unicast(u'protocol error', self)
            elif params[0] == u'get':
                if len(params) != 2:
                    ws.unicast(u'protocol error', self)
                elif params[1] in mars_data.SUPPORTED_DATA:
                    data = mars_data.get_data(params[1])
                    ws.send('server:%s:%s' % (params[1], data), self)
                else:
                    ws.unicast(u'protocol error', self)
            elif params[0] == u'control':
                ws.request_control(self)
            elif params[0] == u'admin':
                if len(params) != 2:
                    ws.unicast(u'protocol error', self)
                else:
                    self.admin = ws.request_admin(self, params[1])
            elif params[0] == u'send':
                if ws.current is not self and self.admin:
                    ws.request_control(self)
                if ws.current is self:
                    raspberry.send(params[1].encode('utf-8')+'\n')
                else:
                    self.fails += 1
            elif params[0] == u'ac':
                if len(params) < 2:
                    ws.unicast(u'protocol error', self)
                elif self.admin:
                    self.admin_control(params[1:])
                else:
                    ws.unicast(u'admin needed', self)
            elif params[0] == u'pong':
                if len(params) < 2:
                    ws.unicast(u'protocol error', self)
                elif params[1] not in self.pings:
                    ws.unicast(u'bad ping', self)
                else:
                    now = ioloop.time()
                    latency = now-self.pings.pop(params[1])
                    ws.send('server:pingtime:%d' % (latency*1000,))
        except (IOError, mars_data.pickle.PickleError) as error:
            ws.unicast_error(error, self)
    
    def send_ping(self):
        '''
        Send a ping request to client to see if they are Still Alive.
        '''
        now = ioloop.time()
        pingid = '%0.2f'%(now,)
        self.pings[pingid] = now
        ws.send('server:ping:%s' % (pingid,), self)
        if len(self.pings) > 2:
            try:
                self.close()
            except AttributeError:
                ws.unregister_client(self)
        else:
            ioloop.add_timeout(self.PING_INTERVAL, self.send_ping)
        
    def on_close(self):
        '''
        Called by Tornado when connection is closed.
        '''
        ws.unregister_client(self)
    
    def admin_control(self, params):
        '''
        Administrator control functions.
        '''
        if params[0] == u'halt':
            raspberry.halt()
            ws.broadcast(u'halt')
        elif params[0] == u'unhalt':
            raspberry.unhalt()
            ws.broadcast(u'unhalt')
        elif params[0] == u'get clients':
            now = datetime.datetime.utcnow()
            if ws.current is not None:
                ws.send('server:current:%x' % (id(ws.current),), self)
            else:
                ws.send('server:current:', self)
            for client in ws._clients:
                connected = now - client.connected
                ws.send('server:client:%x:%d:%d' % (
                    id(client), connected.days, connected.seconds,
                ), self)
            for client in ws._admins:
                connected = now - client.connected
                ws.send('server:admin:%x:%d:%d' % (
                    id(client), connected.days, connected.seconds,
                ), self)
        elif params[0] == u'set control':
            if len(params) != 2:
                ws.unicast(u'protocol error', self)
            else:
                for client in ws._clients|ws._admins:
                    if params[1] == '%x'%id(client):
                        ws.change_control(client)
                        break
                else:
                    ws.unicast(u'client not found', self)
        elif params[0] == u'reconnect':
            raspberry.close()


class StaticFileHandler(tornado.web.StaticFileHandler):
    def get_cache_time(self, path, modified, mime_type):
        return 0


class MarsVideoFeed(tornado.web.RequestHandler):
    '''
    Handler for requests of video data from web browsers.
    '''
    MAX_FAILS = 100
    def compute_etag(self):
        return None
    
    @tornado.web.asynchronous
    def get(self):
        '''
        GET /video
        '''
        print "*** REQUEST VIDEO ***"
        self.ready = True
        self.fails = 0
        print str(hash(self))
        raspberry.add_video_feed(self)
        self.set_status(200)
        self.set_header('Content-Type',
            'multipart/x-mixed-replace;boundary=boundarydonotcross')

    def got_video(self, data):
        '''
        Got some video data.
        '''
        # print "Sending video"
        if not self.ready:
            self.fails += 1
            if self.fails > self.MAX_FAILS:
                print "Forcing connection close"
                self.finish()
            return
        else:
            self.ready = False
            self.fails = 0
        self.write('--boundarydonotcross\r\n')
        self.write('Content-Type: image/jpeg\r\n')
        self.write('Content-Length: %d\r\n\r\n' % len(data))
        self.write(data)
        self.write('\r\n')
        self.flush(callback=self.set_ready)
    
    def set_ready(self):
        '''
        Called by tornado when all data has been sent.
        '''
        self.ready = True
    
    def on_finish(self):
        '''
        Connection closed.
        '''
        print "*** END VIDEO ***"
        print str(hash(self))
        raspberry.remove_video_feed(self)


application = tornado.web.Application([
    (r'/mars', MarsWebSocket),
    (r'/video', MarsVideoFeed),
    (r'/(.*)', StaticFileHandler, {'path': 'http'}),
])

if __name__ == '__main__':
    ioloop = tornado.ioloop.IOLoop.instance()
    application.listen(8888, io_loop=ioloop)
    raspberry = RPiConnection(ioloop)
    print "Server start"
    ioloop.start()

